import random
import asyncio
import logging
import numpy as np
import traceback

from rpcudp.protocol import RPCProtocol

from kademlia.node import Node
from kademlia.routing import RoutingTable
from kademlia.utils import digest, tensor_to_hex
from kademlia.crawling import NodeSpiderCrawl

import ipfshttpclient

from watermarking.dct_watermark import DCT_Watermark
from watermarking.utils import (
    make_qrcode,
    PIL_to_cv2,
)
import imagehash
import hashlib
from PIL import Image
import cv2
import math
import yaml


log = logging.getLogger(__name__)  # pylint: disable=invalid-name
log.setLevel(logging.WARNING)

imgs = {}
config = yaml.load(
    open("/mnt/e/pydev/drpchain/config.yaml").read(), Loader=yaml.FullLoader
)


class KademliaProtocol(RPCProtocol):
    def __init__(self, source_node, storage, ksize):
        # timeout
        RPCProtocol.__init__(self, 5)
        self.router = RoutingTable(self, ksize, source_node)
        self.storage = storage
        self.source_node = source_node
        self.contract = None
        self.event_filter = None
        self.ipfs_client = ipfshttpclient.connect(config["data"]["ipfs"])

    def get_refresh_ids(self):
        """
        Get ids to search for to keep old buckets up to date.
        """
        ids = []
        for bucket in self.router.lonely_buckets():
            rid = random.randint(*bucket.range).to_bytes(20, byteorder="big")
            ids.append(rid)
        return ids

    def rpc_stun(self, sender):  # pylint: disable=no-self-use
        return sender

    def rpc_ping(self, sender, nodeid):
        source = Node(nodeid, sender[0], sender[1])
        self.welcome_if_new(source)
        return self.source_node.id

    def rpc_store(self, sender, nodeid, key, value):
        try:
            source = Node(nodeid, sender[0], sender[1])
            self.welcome_if_new(source)
            if isinstance(key, list):
                for kv in zip(key, value):
                    self.storage.set(kv[0], kv[1])
            # else:
            #     self.storage.set(key, value)
            log.debug(
                "got a store request from %s, storing '%s'='%s'",
                sender,
                str(key),
                str(value),
            )
        except:
            print("rpc_store except:", sender, nodeid, key, value)

        return True

    def rpc_find_node(self, sender, nodeid, key):
        log.info("finding neighbors of %i in local table", int(nodeid.hex(), 16))
        source = Node(nodeid, sender[0], sender[1])
        self.welcome_if_new(source)
        node = Node(key)
        neighbors = self.router.find_neighbors(node, exclude=source)
        return list(map(tuple, neighbors))

    def rpc_find_value(self, sender, nodeid, key):
        source = Node(nodeid, sender[0], sender[1])
        self.welcome_if_new(source)
        value = self.storage.get(key, None)
        if value is None:
            return self.rpc_find_node(sender, nodeid, key)
        return {"value": value}

    # https://pypi.org/project/ImageHash/
    # https://ieeexplore.ieee.org/document/6980335
    async def rpc_build_block(self, sender, nodeid, key, value):
        try:
            # source = Node(nodeid, sender[0], sender[1])
            raw_bytes = self.ipfs_client.cat(value["rawipfsloc"])
            raw_img = Image.frombytes("RGB", value["rawsize"], raw_bytes)
            # raw_img.save(f"{result_dir}/a_{value['title']}")
            raw_phash = imagehash.phash(raw_img)
            # print(str(raw_phash))

            while raw_img.size[0] < 375 or raw_img.size[1] < 375:
                raw_img = raw_img.resize(
                    (int(raw_img.size[0] * 1.2), int(raw_img.size[1] * 1.2)), 1
                )

            # print(
            #     "rpc_build_block",
            #     nodeid.hex(),
            #     raw_img.size,
            #     raw_img.size[0] / raw_img.size[1],
            # )
            
            token_sig = [
                str(raw_phash),
                value["author"],
                value["title"],
                value["timestamp"],
            ]

            raw_img = cv2.cvtColor(np.array(raw_img), cv2.COLOR_RGB2BGR)
            watermark = make_qrcode(token_sig)
            dct = DCT_Watermark()

            wmd = dct.embed(raw_img, watermark)

            # Convert to PIL and Store in IPFS
            wmd = cv2.cvtColor(wmd, cv2.COLOR_BGR2RGB)
            wmd = Image.fromarray(wmd, "RGB")

            wmd_bytes = wmd.tobytes()
            wmd_phash = imagehash.phash(wmd)
            wmd_ipfs_loc = self.ipfs_client.add_bytes(wmd_bytes)
            min_hd_loc, min_hd = self.storage.hamming_distances(
                wmd_phash.hash.flatten()
            )

            log_proof = (
                hashlib.sha1(
                    f"{min_hd_loc}{list(self.storage.data.keys())[0]}".encode("UTF-8")
                )
                .digest()
                .hex()
            )

            # build candidate block broadcast
            wmd_crypto_hash = digest(wmd_bytes)
            candidate_block = {
                "wmdphash": str(wmd_phash),
                "wmdcryptohash": wmd_crypto_hash.hex(),
                "tokensig": token_sig,
                "wmdipfsloc": wmd_ipfs_loc,
                "wmdsize": wmd.size,
                "provider": [nodeid, *sender],
                "logindex": len(self.storage.data),
                "logproof": log_proof,
            }

            # wmd image crypto hash
            node = Node(wmd_crypto_hash)
            nearest = self.router.find_neighbors(node, k=30)
            if not nearest:
                log.warning(
                    "There are no known neighbors to set key %s", wmd_crypto_hash.hex()
                )
                return False
            spider = NodeSpiderCrawl(self, node, nearest, 30, 1)
            nodes = await spider.find()

            validation_results = [
                self.call_validate_block(n, wmd_crypto_hash, candidate_block)
                for n in nodes
            ]
            # collect block validation result
            validation_responses = await asyncio.gather(*validation_results)
            store_responses = []
            if validation_responses != None and len(validation_responses) > 0:
                # if vote in favor r[1][0] is True
                vote_in_favor = np.count_nonzero(
                    np.array(
                        [r[1][0] if r != None else False for r in validation_responses]
                    )
                )
                if vote_in_favor >= math.ceil((len(nearest) - 1) * 0.66667):
                    # update local ledger
                    self.storage.set(wmd_ipfs_loc, str(wmd_phash))

                    # log consistency detect and update
                    log_counts = np.array([r[1][4] for r in validation_responses])
                    log_diff = len(self.storage.data) - log_counts

                    # ledger broadcast
                    results = []
                    for i in range(len(nodes)):
                        update_keys = list(self.storage.data.keys())[-log_diff[i] :]
                        update_vals = list(self.storage.data.values())[-log_diff[i] :]
                        results.append(
                            self.call_store(nodes[i], update_keys, update_vals)
                        )

                    store_responses = await asyncio.gather(*results)
                    return vote_in_favor
                else:
                    reason = np.array(
                        [
                            r[1][1] if r != None and r[1][0] == False else None
                            for r in validation_responses
                        ]
                    )

                    reason = reason[reason != None]
                    return vote_in_favor, str(reason[0])

        except:
            traceback.print_exc()
            return False

    async def rpc_validate_block(self, sender, nodeid, key, value):
        try:
            node = Node(key)
            # source = Node(nodeid, sender[0], sender[1])
            # enable muilt-leader count
            nearest = self.router.find_neighbors(node, k=30)

            # Leader legitimacy
            # xor_notin = f"{sender[0]}:{sender[1]}" not in [str(n) for n in nearest]
            # log_index_notin = len(self.storage.data) > value["logindex"]
            # if log_index_notin:
            #     return (
            #         False,
            #         "Leader legitimacy",
            #         value["wmdcryptohash"],
            #         f"{self.source_node.id.hex()}:{str(self.source_node)}",
            #         len(self.storage.data),
            #         log_index_notin,
            #     )

            # Perceptual hash consistency
            wmd_bytes = self.ipfs_client.cat(value["wmdipfsloc"])
            wmd = Image.frombytes("RGB", value["wmdsize"], wmd_bytes)
            # wmd.save(f"{result_dir}/wmd_{value['wmdcryptohash']}.jpg")
            wmd_phash = imagehash.phash(wmd)
            if str(wmd_phash) != value["wmdphash"]:
                return (
                    False,
                    "Perceptual hash consistency",
                    value["wmdcryptohash"],
                    f"{self.source_node.id.hex()}:{str(self.source_node)}",
                    len(self.storage.data),
                )

            # Originality verication
            min_hd_loc = None
            log_proof = None
            min_hd = 0

            if len(self.storage.data) is not 0:
                min_hd_loc, min_hd = self.storage.hamming_distances(
                    wmd_phash.hash.flatten()
                )

                if min_hd < 10:
                    return (
                        False,
                        "Originality verication",
                        value["wmdcryptohash"],
                        f"{self.source_node.id.hex()}:{str(self.source_node)}",
                        len(self.storage.data),
                    )

                log_proof = (
                    hashlib.sha1(
                        f"{min_hd_loc}{list(self.storage.data.keys())[-1]}".encode(
                            "UTF-8"
                        )
                    )
                    .digest()
                    .hex()
                )

                # Leader legitimacy
                # xor_notin = f"{sender[0]}:{sender[1]}" not in [str(n) for n in nearest]
                # log_index_notin = len(self.storage.data) > value["logindex"]
                log_index_notin = value["logproof"] == log_proof
                if log_index_notin:
                    return (
                        False,
                        "Leader legitimacy",
                        value["wmdcryptohash"],
                        f"{self.source_node.id.hex()}:{str(self.source_node)}",
                        len(self.storage.data),
                        log_index_notin,
                    )

            # Watermark verification
            dct = DCT_Watermark()
            qrcode = dct.extract(PIL_to_cv2(wmd, 3))
            qrcode = qrcode.astype(np.uint8)
            detector = cv2.QRCodeDetector()
            data, vertices_array, _ = detector.detectAndDecode(qrcode)
            if vertices_array is None:
                return False, "Watermark verification", value["wmdcryptohash"]

            # Vote in favor, Proof to validate log consistency
            return (
                True,
                value["wmdcryptohash"],
                log_proof,
                f"{self.source_node.id.hex()}:{str(self.source_node)}",
                len(self.storage.data),
            )

        except Exception as e:
            traceback.print_exc()
            return (
                False,
                f"Error:{str(e)}",
                value["wmdcryptohash"],
                f"{self.source_node.id.hex()}:{str(self.source_node)}",
                len(self.storage.data),
            )

    async def call_find_node(self, node_to_ask, node_to_find):
        address = (node_to_ask.ip, node_to_ask.port)
        result = await self.find_node(address, self.source_node.id, node_to_find.id)
        return self.handle_call_response(result, node_to_ask)

    async def call_find_value(self, node_to_ask, node_to_find):
        address = (node_to_ask.ip, node_to_ask.port)
        result = await self.find_value(address, self.source_node.id, node_to_find.id)
        return self.handle_call_response(result, node_to_ask)

    async def call_ping(self, node_to_ask):
        address = (node_to_ask.ip, node_to_ask.port)
        result = await self.ping(address, self.source_node.id)
        return self.handle_call_response(result, node_to_ask)

    async def call_store(self, node_to_ask, key, value):
        address = (node_to_ask.ip, node_to_ask.port)
        result = await self.store(address, self.source_node.id, key, value)
        return self.handle_call_response(result, node_to_ask)

    #### Add by Drpchain Protocol
    async def call_build_block(self, node_to_ask, key, value):
        address = (node_to_ask.ip, node_to_ask.port)
        result = await self.build_block(address, self.source_node.id, key, value)
        return self.handle_call_response(result, node_to_ask)

    async def call_validate_block(self, node_to_ask, key, value):
        address = (node_to_ask.ip, node_to_ask.port)
        result = await self.validate_block(address, self.source_node.id, key, value)
        return self.handle_call_response(result, node_to_ask)

    def welcome_if_new(self, node):
        """
        Given a new node, send it all the keys/values it should be storing,
        then add it to the routing table.

        @param node: A new node that just joined (or that we just found out
        about).

        Process:
        For each key in storage, get k closest nodes.  If newnode is closer
        than the furtherst in that list, and the node for this server
        is closer than the closest in that list, then store the key/value
        on the new node (per section 2.5 of the paper)
        """
        if not self.router.is_new_node(node):
            return

        log.info("never seen %s before, adding to router", node)
        for key, value in self.storage:
            keynode = Node(digest(key))
            neighbors = self.router.find_neighbors(keynode)
            if neighbors:
                last = neighbors[-1].distance_to(keynode)
                new_node_close = node.distance_to(keynode) < last
                first = neighbors[0].distance_to(keynode)
                this_closest = self.source_node.distance_to(keynode) < first
            if not neighbors or (new_node_close and this_closest):
                asyncio.ensure_future(self.call_store(node, key, value))
        self.router.add_contact(node)

    def handle_call_response(self, result, node):
        """
        If we get a response, add the node to the routing table.  If
        we get no response, make sure it's removed from the routing table.
        """
        if not result[0]:
            log.warning("no response from %s, removing from router", node)
            self.router.remove_contact(node)
            return result

        log.info("got successful response from %s", node)
        self.welcome_if_new(node)
        return result
